package myLock;

import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.locks.AbstractQueuedSynchronizer;
import java.util.concurrent.locks.LockSupport;

/**
 * 手写实现AQS，需要功能点
 * 1. 加锁 lock
 *      1.1
 * 2. 解锁 unlock
 * 3.
 */
public class MyLock {
    // 锁的状态（可重入）
    private final AtomicInteger state = new AtomicInteger(0); // 只保留一个权威的状态源：AtomicInteger
    // 锁的主人
    private Thread owner = null;
    // AQS头节点
    private AtomicReference<Node> head = new AtomicReference<>(new Node());
    // AQS尾节点
    private AtomicReference<Node> tail = new AtomicReference<>(head.get());

    public void lock() {
        // 非公平，先尝试直接获取锁
        if(state.compareAndSet(0, 1)) {
            owner = Thread.currentThread();
            return;
        }
        // 处理重入
        if (owner == Thread.currentThread()) {
            state.incrementAndGet();
            return;
        }
        // 获取锁失败，加入队列
        Node node = addWaiter();
        acquireQueued(node);
    }

    /**
     * 将当前线程加入等待队列的尾部
     * @return 当前线程的节点
     */
    private Node addWaiter() {
        Node newNode = new Node(Thread.currentThread());
        // 快速尝试在尾部加入
        Node pred = tail.get();
        if (pred != null) {
            newNode.prev = pred;
            if (tail.compareAndSet(pred, newNode)) {
                pred.next = newNode;
                return newNode;
            }
        }
        // 如果队列为空或CAS失败，进入完整的入队逻辑
        enq(newNode);
        return newNode;
    }

    /**
     * 通过自旋CAS保证节点成功入队
     * @param node 要入队的节点
     */
    private void enq(final Node node) {
       for (;;) {
           Node t = tail.get();
           if (t == null) { // 队列未初始化
               // 初始化head和tail为一个虚拟节点
               if (head.compareAndSet(null, new Node())) {
                   tail.set(head.get());
               }
           }
           else {
               node.prev = t;
               if (tail.compareAndSet(t, node)) {
                   t.next = node; // 链接成功，可以退出了
                   return;
               }
           }
       }
    }

    /**
     * 在队列中等待并获取锁
     * @param node 当前线程的节点
     */
    private void acquireQueued(final Node node) {
        boolean failed = false;
        try {
            for (;;) {
                final Node p = node.prev;
                // 如果前驱是头节点，说明轮到自己了，再次尝试获取锁
                if (p == head.get() && state.compareAndSet(0, 1)) {
                    setHead(node); // 获取成功，将自己设为头节点
                    p.next = null; // help GC
                    failed = false;
                    owner = Thread.currentThread();
                    return;
                }
                // 如果没轮到自己，或者轮到自己但获取锁失败，判断是否应该阻塞
                if (shouldParkAfterFailedAcquire(p, node)) {
                    LockSupport.park(this); // 阻塞当前线程
                }
            }
        } finally {
            if (failed) {
                // 如果因为异常等原因退出，需要取消节点
                cancelAcquire(node);
            }
        }
    }

    private void setHead(Node node) {
        head.set(node);
        node.thread = null; // 头节点不需要线程
        node.prev = null;   // help GC
    }

    private static boolean shouldParkAfterFailedAcquire(Node pred, Node node) {
        int ws = pred.waitStatus;
        if (ws == Node.SIGNAL) {
            // 前驱节点状态是SIGNAL，说明它承诺会唤醒我，我可以安心阻塞了
            return true;
        }
        if (ws > 0) {
            // 前驱节点已取消(CANCELLED)，需要向前遍历，跳过所有已取消的节点
            do {
                node.prev = pred = pred.prev;
            } while (pred.waitStatus > 0);
            pred.next = node;
        }
        else {
            // 前驱节点是 0 或 PROPAGATE，用CAS把它设置为SIGNAL
            // 期望它下次释放锁时能唤醒我
            pred.compareAndSetWaitStatus(ws, Node.SIGNAL);
        }
        return false; // 本次不阻塞，循环下一次会再次检查
    }

    public void unlock() {
        if (owner != Thread.currentThread()) {
            throw new IllegalMonitorStateException();
        }

        // 原子递减并获取新值
        int current_state = state.decrementAndGet();
        if (current_state == 0) { // 完成释放锁
            owner = null;
            Node h = head.get();
            // 如果队列不为空，且后继节点需要被唤醒
            if (h != null && h.waitStatus != 0) {
                unparkSuccessor(h);
            }
        }
        // 如果 current_state > 0，说明只是释放了一次重入锁，不需要做任何其他事
    }

    /**
     * 唤醒后继节点
     * @param node 当前节点（通常是头节点）
     */
    private void unparkSuccessor(Node node) {
        int ws = node.waitStatus;
        if (ws < 0) {
            // 清理自己的状态，允许失败
            node.compareAndSetWaitStatus(ws, 0);
        }
        Node s = node.next;
        // 如果后继节点是null或者已取消，从尾部向前遍历找到最靠前的有效等待者
        if (s == null || s.waitStatus > 0) {
            s = null;
            for (Node t = tail.get(); t != null && t != node; t = t.prev) {
                // 包含 t.waitStatus == 0 是为了处理一个时间窗口：一个节点已经成功加入队列，但它的线程还未来得及将它的前驱节点状态设置为 SIGNAL
                if (t.waitStatus <= 0) {
                    s = t;
                }
            }
        }
        if (s != null) {
            LockSupport.unpark(s.thread);
        }
    }

    private void cancelAcquire(Node node) {
        // 实现省略，主要逻辑是标记节点为CANCELLED，并尝试将其从链中断开
        node.waitStatus = Node.CANCELLED;
        // ... 复杂的断链和唤醒后继者逻辑
    }

    class Node {
        static final int CANCELLED = 1;
        static final int SIGNAL = -1;

        volatile Node prev;
        volatile Node next;
        volatile Thread thread;
        volatile int waitStatus;

        Node() { }

        Node(Thread thread) {
            this.thread = thread;
        }

        // 使用CAS来原子更新waitStatus
        final boolean compareAndSetWaitStatus(int expect, int update) {
            // 实际AQS使用Unsafe.compareAndSwapInt，这里我们模拟一下
            // 简单起见，这里直接赋值，但在真实并发场景下必须是原子操作
            // 为了演示，我们假设这个操作是原子的
            if (this.waitStatus == expect) {
                this.waitStatus = update;
                return true;
            }
            return false;
        }
    }
}
